(module type-unification mzscheme
  (require "types.scm")
  (require (lib "plt-match.ss"))
  (require (only (lib "1.ss" "srfi") append-reverse!))
  
  ;; This structure represents type constraints
  (define-struct constraint (lhs rhs))
  
  (provide make-constraint)
  
  ;; Adds a substitution id |-> ty to the environment.
  ;; it will also change all substitutions of the form id' |-> id
  ;; to be id' |-> ty
  (define (add-subst id ty env)
    (hash-table-for-each env
     (lambda (k v)
       (if (equal? v id)
           (hash-table-put! env k ty))))
    (hash-table-put! env id ty))
  
  (define ((chase env) id)
    (display (format "Chasing ~S -- in hash table ~S~%" id (hash-table-get env id #f)))
    (cond ((hash-table-get env id #f) => (lambda (x) (get-id env x)))
          (else id)))
  
  (define (get-id env x)
    (and (ty-var? x)
        ((chase env) (ty-var-identifier x))))
  
  ;; Given a list of type constraints, this returns the environment containing
  ;; satisfying assignments
  ;; Returns either a hash table containing the satisfying type assignments
  ;; or a string describe the conflict
  (define (unify constraints)
    (let ((env (make-hash-table)))
    (let unify-loop ((stack constraints))
      (if (null? stack) env
          (let ((lhs  (constraint-lhs (car stack)))
                (rhs  (constraint-rhs (car stack)))
                (rest (cdr stack)))
            (let ((lhs-id (get-id env lhs))
                  (rhs-id (get-id env rhs)))
              (display (format "Unifying ~S ~%     and ~S ~%   (ids are ~S and ~S)~%"
                               lhs rhs lhs-id rhs-id))
              (cond 
                ((and lhs-id rhs-id (eq? lhs-id rhs-id)) (unify-loop rest))
                (lhs-id 
                 (add-subst lhs-id rhs env)
                 (unify-loop rest))
                (rhs-id 
                 (add-subst rhs-id lhs env) 
                 (unify-loop rest))
                ((and (list? lhs) (list? rhs))
                 (if (= (length lhs) (length rhs))
                     (unify-loop (append-reverse! (map make-constraint lhs rhs) rest))
                     (raise-syntax-error #f "Unification failed because lists are of different length" lhs rhs)))
                ((and (ty-app? lhs) (ty-app? rhs))
                 (if (= (length (ty-app-args lhs)) (length (ty-app-args rhs)))
                     (unify-loop (append-reverse! (map make-constraint 
                                                       (ty-app-args lhs)
                                                       (ty-app-args rhs))
                                                  rest))
                     (raise-syntax-error #f "Unification failed because arg lists are of different length" lhs rhs)))
                (else (raise-syntax-error #f "Unification failed" lhs rhs)))))))))
  
  (define ((substitute env) ty)
    (match ty
      ((struct ty-var (id))
       (hash-table-get env id ty))
      ((struct ty-app (con args))
       (make-ty-app con (map (substitute env) args)))))
  
  (define (compose func1 func2)
    (let ((func1-args (get-arrow-args (instantiate func1)))
          (func2-args (get-arrow-args (instantiate func2))))
      (if (or (null? func1-args) (null? func2-args))
          (raise-syntax-error #f "Composition expects two arrows as arguments" func1 func2)
          (let ((func1-argty (car  func1-args))
                (func1-resty (cadr func1-args))
                (func2-argty (car  func2-args))
                (func2-resty (cadr func2-args)))
            (let ((env (unify (list (make-constraint func1-resty func2-argty)))))
              (make-ty-app (make-tycon-pure-arrow) (list ((substitute env) func1-argty)
                                                         ((substitute env) func2-resty))))))))
  (provide compose substitute unify)
)
    
  